# -*- coding: utf-8 -*-
"""
Created on 2021/7/30 0:00

@author: 张径舟
"""
import os
from io import BytesIO

import mxnet as mx
from mxnet import image, gpu
import gluoncv
from gluoncv.data.transforms.presets.segmentation import test_transform
from gluoncv.utils.viz import get_color_pallete, plot_image
from mxnet.gluon.data import dataset
from mxnet.gluon.data import DataLoader
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np
import pandas as pd
import cv2
from PIL import Image
import sqlite3

from mxnet.gluon.data.vision import transforms
import warnings

warnings.filterwarnings(action='once')
warnings.filterwarnings("ignore")
col_map = {0: 'road', 1: 'sidewalk', 2: 'building', 3: 'wall', 4: 'fence', 5: 'pole', 6: 'traffic light',
           7: 'traffic sign', 8: 'vegetation', 9: 'terrain', 10: 'sky', 11: 'person', 12: 'rider',
           13: 'car', 14: 'truck', 15: 'bus', 16: 'train', 17: 'motorcycle', 18: 'bicycle'}
BATCH_SIZE = 8

db_path = r"C:\spider\streetview0726.db"
result_db_path = r"C:\spider\streetview_segmentation_result.db"


def TRANSFORM(img, ctx):
    transform_fn = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([.485, .456, .406], [.229, .224, .225])
    ])
    img = transform_fn(img)
    img = img.as_in_context(ctx)
    return img


def load_img(binary):
    img = Image.open(BytesIO(binary))
    img = np.array(img, dtype=np.dtype('i4'))
    return img


def Predict(db_path, result_db_path, model, ctx):
    import time
    # ds = StreetViewPredictDataset(IMG_DIR, ctx, transform=TRANSFORM)
    # dl = DataLoader(ds, BATCH_SIZE, shuffle=False, last_batch='keep')
    print("start")
    classes = ['road', 'sidewalk', 'building', 'wall', 'fence',
               'pole', 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky',
               'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', 'bicycle']
    #df = pd.DataFrame(columns=classes)
    result_db = sqlite3.connect(result_db_path)
    sql = f"""
    CREATE TABLE IF NOT EXISTS result 
    (
    img_id TEXT, 
    heading FLOAT, 
    mask BLOB,
    {','.join([curr_class + " FLOAT" if 'traffic' not in curr_class else '_'.join(curr_class.split()) + " FLOAT" for curr_class in classes])}
    )
    """
    result_cursor = result_db.cursor()
    result_cursor.execute(sql)
    result_cursor.execute("PRAGMA INDEX_LIST(result)")
    result_db_indices = result_cursor.fetchall()
    result_db_indices = [index[1] for index in result_db_indices]
    if "img_id_index" not in result_db_indices:
        result_cursor.execute("CREATE INDEX img_id_index on result (img_id)")
    if "bd09_index" not in result_db_indices:
        result_cursor.execute("CREATE INDEX heading_index on result (heading)")

    db = sqlite3.connect(db_path, uri=True)
    cursor = db.cursor()
    total = cursor.execute("SELECT COUNT(*) FROM street_view").fetchone()[0]
    cursor.arraysize = BATCH_SIZE
    cursor.execute("SELECT img_id, image, heading FROM street_view")
    count = 0
    lines = cursor.fetchmany()
    aaa = time.time()
    while len(lines) != 0:
        bbb = time.time()
        images = []
        for line in lines:
            images.append(load_img(line[1]))
        images = mx.nd.array(images)
        img = TRANSFORM(images, ctx)
        output = model.predict(img)
        predicts = mx.nd.argmax(output, 1).asnumpy()
        for index in range(predicts.shape[0]):
            img_id = lines[index][0]
            heading = lines[index][2]
            predict = predicts[index]
            #base = Image.open(BytesIO(lines[index][1]))
            #base.save(os.path.join(img_id.split(".")[0] + f"_{heading}.png"))
            mask = get_color_pallete(predict, 'citys')
            #mask.save(os.path.join(img_id.split(".")[0] + f"_{lines[index][2]}_mask.png"))
            pred = [(len(predict[predict == i]) / (predict.shape[0] * predict.shape[1])) for i in range(19)]
            #Image.new().convert()
            stream = BytesIO()
            mask.convert('RGB').save(stream, format="JPEG")
            result_cursor.execute(f"INSERT INTO result VALUES (?,?,?,{','.join(['?'] * 19)})", [img_id, heading, stream.getvalue()] + pred)
        count += len(lines)
        if count >= 50:
            result_db.commit()
        print(count, total,
              "已用时: %.6f" % (time.time() - aaa),
              "batch用时: %.6f" % (time.time() - bbb),
              "预计剩余时间: %.6f" % ((time.time() - aaa) / count * (total - count)))
        lines = cursor.fetchmany()
        #break
    result_db.commit()


if __name__ == "__main__":
    ctx = mx.gpu(0)
    model = gluoncv.model_zoo.get_model('psp_resnet101_citys', ctx=ctx, pretrained=True)
    Predict(db_path, result_db_path, model, ctx)
